package com.ruoyi.ai.splitter;

import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.splitter.TextSplitter;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class LineTextSplitter extends TextSplitter {

    public static final String SPLIT_STR = "\\? 答:";
    private int maxLength = Integer.MAX_VALUE; // 默认不限制每行的最大长度

    /**
     * 设置每行的最大长度。如果某行超过此长度，则进一步分割。
     */
    public void setMaxLength(int maxLength) {
        if (maxLength <= 0) {
            throw new IllegalArgumentException("Max length must be greater than 0.");
        }
        this.maxLength = maxLength;
    }

    /**
     * 实现抽象方法 splitText，按行分割文本。
     */
    @Override
    protected List<String> splitText(String text) {
        List<String> lines = new ArrayList<>();
        String[] rawLines = text.split("\\r?\\n"); // 按换行符分割文本

        for (String line : rawLines) {
            line = line.trim(); // 去除首尾空白字符
            if (!line.isEmpty()) { // 跳过空行
                lines.addAll(splitLineByMaxLength(line)); // 如果需要，按最大长度进一步分割
            }
        }

        return lines;
    }

    /**
     * 按最大长度分割单行文本。
     */
    private List<String> splitLineByMaxLength(String line) {
        List<String> chunks = new ArrayList<>();
        if (line.length() <= maxLength) {
            chunks.add(line); // 如果行长度小于等于最大长度，直接添加
        } else {
            // 分割长行
            int start = 0;
            while (start < line.length()) {
                int end = Math.min(start + maxLength, line.length());
                chunks.add(line.substring(start, end));
                start += maxLength;
            }
        }
        return chunks;
    }

    @Override
    public List<Document> split(Document document) {
        String text = document.getText();
        List<String> splitText = this.splitText(text);
        if (splitText.isEmpty()) {
            return List.of(document);
        }
        List<Document> documents = new ArrayList<>();
        for (String chunk : splitText) {
            //解析格式：你叫什么名字? 答:我是雅橙聊天助手。
            String[] split = chunk.split(SPLIT_STR);
            String question = split[0];
            String answer = split[1];
            // 创建metadata，包含标题下的内容
            Map<String, Object> metadata = new HashMap<>(document.getMetadata()); // 保留原始metadata
            metadata.put("content", answer);
            documents.add(new Document(question,metadata));
        }
        return documents;
    }

    /**
     * 测试主程序
     */
    public static void main(String[] args) {
        // 示例文本
        String inputText = """
                你叫什么名字? 答:我是雅橙聊天助手。
                你几岁了啊? 答:我1岁不到。
                """;

        // 创建 LineTextSplitter 实例
        LineTextSplitter splitter = new LineTextSplitter();
        splitter.setMaxLength(50); // 设置每行的最大长度为 50

        // 模拟 Document 输入
        Document document = new Document(inputText, Map.of("source", "example.txt"));

        // 调用 split 方法
        List<Document> result = splitter.split(document);

        // 打印分割结果
        for (int i = 0; i < result.size(); i++) {
            System.out.println("Chunk " + (i + 1) + ": " + result.get(i).getText());
        }
    }
}
